# @matching/single_match.py
import math
import time
from bson import ObjectId
from entity.function_candidate import FunctionCandidate
from config.config import MatchType,MatchMethod
from matching.filter_by_name import get_best_name_match_from_ids
from utils.calculate_utils import build_graph_from_feature
from matching.first_stage import first_stage_filter
from matching.isomorphic_stagev3 import isomorphic_ratio_stage
from utils.mongo_utils import *
import logging
from pymongo import  ASCENDING, DESCENDING


# 同构子树的语义相似度阈值，一般不需要改动
token_similarity_threshold = 0.6
# 同构子树的最终阈值
sameAST_similarity_threshold = 0.5


def find_best_match_for_target(target_id,
     matchType,ArkTS_Library = None, ):
    """
    :param target_id: 待检测的函数的函数id
    :param matchType: 当前匹配的过程类型，要求为TPL或TPL Version（MatchType中枚举）
    :param ArkTS_Library：当匹配类型为TPLVersion时候需要给定已经得到的ArkTS库
       主函数：为待检测的函数找到最佳候选者匹配结果。
       TPL匹配流程中 target为 arkts ,match 为 NPM ,
       但是在TPL version过程中 target 为 npm ，但是 match 为 arkts
       流程：初始化 -> 第一轮筛选 -> 处理结果 ->第二轮筛选 -> 保存结果
       """

    start_time = time.time()
    # --- 1. 初始化阶段 ---
    # 期待的结果集合，主要看与谁对比，TPL就和NPM对比，TPL Version就和ArkTS对比
    match_functions = get_npm_collection() if matchType == MatchType.TPL_MATCH else get_arkts_collection()
    # add_fixed_indexes(match_functions)
    # 待检测函数的集合，方便查询其他属性
    target_functions = get_arkts_collection() if matchType == MatchType.TPL_MATCH else get_npm_collection()
    target_doc, target_name, G1, target_node_count = _initialize_match(target_id, target_functions)
    #判断是否为匿名函数，如果是非匿名函数就先试用名称进行搜索，如果是匿名函数直接使用向量进行搜索。
    is_non_anonymous = bool(target_name and target_name.lower() != 'anonymous')

    # --- 2. 初步筛选，获取到中高候选者 ---
    high_similarity_candidates = [] # 高相似度候选者
    medium_similarity_candidates= [] # 中等相似度候选者
    has_vector = False
    if matchType == MatchType.TPL_MATCH:
        high_similarity_candidates, medium_similarity_candidates, has_vector = first_stage_filter(
            target_doc, target_id , matchType
        )
    elif matchType==MatchType.TPL_VERSION_MATCH:
        if ArkTS_Library is None:
            # 可以选择抛出异常或设置默认值
            raise ValueError("TPL_VERSION_MATCH 模式需要提供 ArkTS_Library 参数")
        high_similarity_candidates, medium_similarity_candidates, has_vector = first_stage_filter(
            target_doc, target_id , matchType, ark_ts_library=ArkTS_Library
        )
    # 检查筛选结果，寻找最佳和最好的一批候选者
    if high_similarity_candidates:
        # 高候选者直接返回就可以
        best_match = high_similarity_candidates[0]

        if is_non_anonymous:
            # 提取所有高相似度候选的ID
            candidate_ids = [ObjectId(c.candidate_id)  for c in high_similarity_candidates]
            # 找出名称最匹配的ID
            best_name_match_id = get_best_name_match_from_ids(
                matchType,candidate_ids, target_name)
            if best_name_match_id:
                # 找到名称匹配的候选
                best_match = next((c for c in high_similarity_candidates if c.candidate_id == best_name_match_id), None)

        all_matches = high_similarity_candidates[:5]
        logging.debug(
            f"找到 {MatchMethod.VECTOR_ONLY.value} {len(all_matches)}个向量候选者,最佳匹配: {"ArkTS" if matchType == MatchType.TPL_VERSION_MATCH else "NPM"}函数 {best_match.candidate_id}, 相似度: {best_match.vector_similarity}")
        processing_time01 = time.time() - start_time
        best_match.match_method = MatchMethod.VECTOR_ONLY.value
        best_match.processing_time = processing_time01
        for match in all_matches:
            match.match_method = MatchMethod.VECTOR_ONLY.value
            match.processing_time = processing_time01
        return best_match,all_matches
    elif medium_similarity_candidates:
        #同构异构的比较
        candidate_ids = [ObjectId(c.candidate_id) for c in medium_similarity_candidates]
        logging.debug(f"开始批量获取 {len(candidate_ids)} 个中等相似度候选函数的完整文档...")
        fetch_query = {
            '_id': {'$in': candidate_ids},
            'ast_feature': {'$exists': True}
        }
        # 从特征库中查找对应候选者的文档
        full_candidate_docs_cursor = match_functions.find(fetch_query, {'_id': 1, 'ast_feature': 1, 'name': 1})
        full_candidate_docs_dict = {str(doc['_id']): doc for doc in full_candidate_docs_cursor}

        enriched_medium_candidates = []
        missing_docs_count = 0
        # 获取这些候选者的文档
        for candidate_info in medium_similarity_candidates:
            candidate_info: FunctionCandidate  # ← 变量注解
            full_doc = full_candidate_docs_dict.get(candidate_info.candidate_id)
            if full_doc:
                candidate_info.candidate_doc = full_doc
                enriched_medium_candidates.append(candidate_info)
            else:
                missing_docs_count += 1
                logging.warning(f"警告：未能获取到候选函数 {candidate_info.candidate_id} 的完整文档")
        if missing_docs_count > 0:
            logging.warning(f"共有 {missing_docs_count} 个候选函数未能获取完整文档")
        # 进行同构异构分析
        isomorphic_candidates = isomorphic_ratio_stage(
            G1, target_node_count,
            enriched_medium_candidates,
            has_vector, "vector_similarity",
            sameAST_similarity_threshold
        )
        isomorphic_candidates.sort(key=lambda x: x.distance)
        above_threshold_matches = [m for m in isomorphic_candidates if
                                   m.isomorphic_ratio >= sameAST_similarity_threshold]
        processing_time = time.time() - start_time
        logging.debug(f"同构、异构比例计算耗时 : {processing_time:.2f}秒")
        # 同构比例大于同构阈值，直接作为候选者
        if above_threshold_matches:
            best_match = above_threshold_matches[0]
            all_matches = isomorphic_candidates[:10] if isomorphic_candidates[0].isomorphic_ratio > 0 else isomorphic_candidates
        else:
            # 同构比例都小于同构阈值，拿distance最小的几个，并用name做一层保险
            if isomorphic_candidates:
                best_match = min(isomorphic_candidates, key=lambda x: x.distance)
                all_matches = isomorphic_candidates[:10]
            else:
                best_match = None
                all_matches = []
        # 和上面一样，用Name做一层保险
        if is_non_anonymous:
            candidate_ids = [ObjectId(c.candidate_id) for c in all_matches]
            best_name_match_id = get_best_name_match_from_ids( matchType, candidate_ids, target_name)
            if best_name_match_id:
                best_match = next((c for c in high_similarity_candidates if c.candidate_id == best_name_match_id),best_match)

        # --- 返回结果 ---
        if best_match:
            best_match: FunctionCandidate
            logging.debug(f"\n最佳匹配: 函数 {best_match.candidate_id}")
            if not math.isnan(best_match.vector_similarity):
                logging.debug(f"向量相似度: {best_match.vector_similarity:.4f}")
            if not math.isnan(
                    best_match.isomorphic_ratio) and best_match.isomorphic_ratio >= sameAST_similarity_threshold:
                logging.debug(f"同构比例: {best_match.isomorphic_ratio:.2%}")
            else:
                logging.debug(f"同构比例: {best_match.isomorphic_ratio:.2%}")
                if best_match.diff_similarity >= 0:
                    logging.debug(f"异构相似度: {best_match.diff_similarity:.4f}")
                if best_match.distance >= 0:
                    logging.debug(f"距离: {best_match.distance:.4f}")
            logging.debug(f"处理完成，总耗时: {processing_time:.2f}秒")

            # 返回最佳匹配ID和相似度
            best_match.match_method = MatchMethod.HYBRID.value
            best_match.processing_time = processing_time
            for match in all_matches:
                match.match_method = MatchMethod.HYBRID.value
                match.processing_time = processing_time

            return best_match, all_matches

        else:
            logging.warning(f"\n同构异构方法未找到匹配，总耗时: {processing_time:.2f}秒")
            return None, 0
    else:
        logging.warning(f"警告：待检测函数 {target_id} 没有函数匹配结果")
        return None

def _initialize_match(target_id, target_functions):
    target_doc = target_functions.find_one({"_id": ObjectId(target_id), "ast_feature": {"$exists": True}})
    if not target_doc:
        raise ValueError(f"找不到包含 特征 的待检测的函数: {target_id}")
    target_name = target_doc.get('name', '').strip()
    function_node_count = target_doc['ast_feature'].get('nodeCount', 0)
    G1 = build_graph_from_feature(target_doc['ast_feature'])
    logging.debug(f"开始为 函数 {target_id} ({target_name}) 寻找匹配。")
    return target_doc, target_name, G1, function_node_count



